import torch
import torch.nn as nn
from PIL import Image
import torch.utils.data as data
import torchvision.transforms as transform
from PIL import Image
import os
import numpy as np

class ImageDataset(data.Dataset):
    def __init__(self, img_source, mask_source, img_path, transform, target_transform):
        self.img_source = img_source # file list
        self.transform = transform
        self.img_path = img_path # img total path
        self.mask_source = mask_source
        self.target_transfrom = target_transform
        self.img_path_list = [] # save individual img path
        self.mask_path_list = [] # save individual mask path
        img_count = 0
        tag_count = 0

        with open(self.img_source, 'r') as f:
            file_list = f.readlines()
            for file in file_list:
                pass

        with open(self.mask_source, 'r') as f:
            file_list_m = f.readlines()
            for file in file_list_m:
                pass

    def __getitem__(self, index):
        img_path = self.img_path_list[index]
        msk_path = self.mask_path_list[index]

        img = Image.open(img_path).convert('RGB')
        msk = Image.open(msk_path).convert('L')

        if self.transform:
            img = self.transform(img)

        return img, msk

    def __len__(self):
        return len(self.img_path_list)







